#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
import traceback
from fastapi.exceptions import HTTPException
from app.core.exceptions.http_exceptions import CustomException
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio.session import AsyncSession
from sqlalchemy.orm import DeclarativeBase, MappedAsDataclass, sessionmaker
from app.core.pkg.log import logger
from app.core.config import config


class Base(DeclarativeBase, MappedAsDataclass):
    pass


DATABASE_URI = f"{config.get('postgres').get('user')}:{config.get('postgres').get('password')}@{config.get('postgres').get('host')}:{config.get('postgres').get('port')}/{config.get('postgres').get('db')}"
DATABASE_PREFIX = config.get('postgres').get('prefix')
DATABASE_URL = f"{DATABASE_PREFIX}{DATABASE_URI}"

async_engine = create_async_engine(DATABASE_URL,
                                   echo=False,
                                   future=True,
                                   pool_size=200,  # 连接池的初始大小
                                   max_overflow=10,  # 允许的最大额外连接数
                                   pool_timeout=30,  # 从连接池获取连接的超时时间
                                   pool_recycle=3600  # 连接在被回收并重新创建之前可以保持打开的时间)
                                   )

local_session = sessionmaker(bind=async_engine, class_=AsyncSession, expire_on_commit=False)


async def async_get_db() -> AsyncSession:
    async_session = local_session
    async with async_session() as db:
        try:
            yield db
        except Exception as e:
            logger.error(traceback.format_exc())
            await db.rollback()
            if isinstance(e, CustomException):
                raise e
            raise HTTPException(status_code=500, detail='服务异常')
        finally:
            await db.close()
